import os, sys,threading,cv2,time,subprocess,shutil
# from enum import Enum
from pydantic import BaseModel, Field
from typing import Optional,Any,Union
from fastapi import FastAPI,  HTTPException,APIRouter, Depends,Query
# from fastapi.staticfiles import StaticFiles
# import numpy
# from collections import defaultdict,deque
# import torch
# 禁止log
os.environ["PADDLEOCR_DISABLE_AUTO_LOGGING_CONFIG"] = "0"
# print("PADDLEOCR_DISABLE_AUTO_LOGGING_CONFIG:",os.getenv("PADDLEOCR_DISABLE_AUTO_LOGGING_CONFIG"))

try:
    import tool.public_tool
except:
    import public_tool
# 环境依赖
# python -m pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/
# python -m pip install paddleocr
# module_require_list=[["paddleocr"]]
# try:
#     from paddleocr import PaddleOCR
# except ImportError:
#     from tool.api_tool import install_modules
#     install_modules(module_require_list)
#     print("matplotlib工具所需模块已检查完成")
from paddleocr import PaddleOCR
from tool.logging_tool import *
from tool.api_tool import APIRes,verify_api_key,auth_key
from tool.image_tool import parse_image_data

ocr_model=None
def get_model():
    global ocr_model
    if ocr_model is None:
        ocr_model = PaddleOCR(
        # doc_orientation_classify_model_name	文档方向分类模型名称（如用于判断文档是否旋转）
        # doc_orientation_classify_model_dir	文档方向分类模型目录（模型文件路径）
        # doc_unwarping_model_name	文档矫正（去畸变）模型名称
        # doc_unwarping_model_dir	文档矫正模型目录（模型文件路径）
        # text_detection_model_name	文本检测模型名称（如 PP-OCRv4_mobile_det）
        # text_detection_model_dir	文本检测模型目录（模型文件路径）
        # textline_orientation_model_name	文本行方向分类模型名称（如用于检测文本行是否倾斜）
        # textline_orientation_model_dir	文本行方向分类模型目录（模型文件路径）
        # textline_orientation_batch_size	文本行方向分类模型推理批量大小
        # text_recognition_model_name	文本识别模型名称（如 PP-OCRv4_mobile_rec）
        # text_recognition_model_dir	文本识别模型目录（模型文件路径）
        # text_recognition_batch_size	文本识别模型推理批量大小
        # use_doc_orientation_classify	是否启用文档方向分类（True/False）
        # use_doc_unwarping	是否启用文档矫正（True/False）
        # use_textline_orientation	是否启用文本行方向分类（True/False）
        # text_det_limit_side_len	文本检测模型输入图片最大边长限制
        # text_det_limit_type	边长限制类型（如 max/min）
        # text_det_thresh	文本检测像素阈值（概率图大于此值判为文本）
        # text_det_box_thresh	文本检测框阈值（检测框平均分数大于此值判为文本区域）
        # text_det_unclip_ratio	文本检测区域扩展系数（扩大检测框面积）
        # text_det_input_shape	文本检测模型输入 shape（如 [3, 640, 640]）
        # text_rec_score_thresh	文本识别分数阈值（识别结果分数大于此值保留）
        # text_rec_input_shape	文本识别模型输入 shape（如 [3, 32, 320]）

            use_doc_orientation_classify=False, # Disables document orientation classification model via this parameter
            use_doc_unwarping=False, # Disables text image rectification model via this parameter
            use_textline_orientation=False, # Disables text line orientation classification model via this parameter
            )
    return ocr_model
class OCRImageArgs(BaseModel):
    image: Optional[str] = Field(default='', description="远程图片,本地图片,base64", examples=["test.jpg","http://oss.580u.com/ai-public/yolov8/test.jpg",  "base64_image_data"])
    # lang: Optional[str] = Field("ch", description="语言类型，默认为中文")
    # use_gpu: Optional[bool] = Field(False, description="是否使用GPU加速，默认为False")
    # use_angle_cls: Optional[bool] = Field(False, description="是否使用文本方向分类器，默认为False")
    # use_unwarp: Optional[bool] = Field(False, description="是否使用文本图像矫正，默认为False")
    debug: Optional[bool] = Field(default=False, description="是否开启调试模式", examples=[True])
    @property
    def frame(self) -> Union[str, Any]:
        """
        获取图片数据
        :return: 图片数据
        """
        frame=None
        try:
            frame = parse_image_data(self.image)
        except:
            warn_print("图片数据读取异常:", self.image)
            frame=None
        if frame is None:
            raise HTTPException(status_code=400, detail="图片数据无效或不存在")
        return frame
    
def start_router()->APIRouter:
    router = APIRouter(prefix="/ocr", tags=["ocr工具"], responses={404: {"description": "/docs"}},dependencies=[Depends(verify_api_key)])
    
    @router.post("/image", summary="图片识别", response_model=APIRes, description="")
    def image(args:OCRImageArgs) -> APIRes:
        info_print("图片识别:")
        if sys.platform=="linux":
            args.debug=False  # linux环境下默认不开启调试模式
        task_id=time_uuid()
        task_dir= os.path.join(tmpdir, '_ocr', f"{task_id}")
        os.makedirs(task_dir, exist_ok=True)
        try:
            frame=args.frame
            if frame is None:
                raise HTTPException(status_code=404, detail=f"图片 {args.image} 不存在")
            task_file=os.path.join(task_dir, "image.jpg")
            cv2.imwrite(task_file, frame)
            info_print("图片保存到:",task_file)
            info_print("图片识别: 开始")
            result = get_model().predict(task_file)
            res_json=[]
            for res in result:
                res.save_to_img(task_dir)
                res.save_to_json(task_dir)
                rec_texts=res.json.get("res", {}).get("rec_texts", [])
                info_print("图片识别: 完成",rec_texts)
                res_json.append(res.json.get("res", {}))
            if args.debug:
                pass
        except Exception as e:
            error_print("图片识别异常:", e)
            raise HTTPException(status_code=500, detail=f"图片识别异常: {str(e)}")
        return APIRes(status_code=200, data=res_json, detail=f"视频流连接失败")
    return router
    # 
def start_app(app: FastAPI):
    get_model()
    app.include_router(start_router())
    info_print("ocr工具路由已启动")
    return app
def stop_app(app: FastAPI):
    info_print("ocr工具路由已停止")
    return app
def test():
    import uvicorn,argparse,sys
    from typing import Union
    from fastapi import FastAPI, Body, HTTPException, Request, Response,APIRouter, Depends
    from fastapi.staticfiles import StaticFiles
    from contextlib import asynccontextmanager
    
    # 定义 lifespan 事件处理器
    @asynccontextmanager
    async def lifespan(app: FastAPI):
        warn_print("等待应用关闭...")
        logging_reset_config()
        yield  # 应用运行期间
        warn_print("应用已经关闭，清理资源...")
        stop_app(app)
    app = FastAPI(lifespan=lifespan,description=f"授权Key:{auth_key()[0]}",)
    start_app(app)
    
    
    uvicorn.run(app, host='0.0.0.0')
    info_print("test end")
    pass
if __name__ == "__main__":
    # 测试模型加载
    # load_local_models()
    test()
    # download_models()
    print("FFMPEG工具加载完成")